Scalable K-FAC Training for Deep Neural Networks With Distributed Preconditioning

The second-order optimization methods, notably the D-KFAC (Distributed Kronecker Factored Approximate Curvature) algorithms, have gained traction on accelerating deep neural network (DNN) training GPU clusters. However, existing algorithms require to compute and communicate a large volume of information, i.e., factors (KFs), before preconditioning gradients, resulting in computation communication overheads as well high memory footprint. In this paper, we propose DP-KFAC, novel distributed scheme that distributes KF constructing tasks at different DNN layers workers. DP-KFAC not only retains convergence property but also enables three benefits: reduced overhead KFs, no low Extensive experiments 64-GPU cluster show reduces by 1.55×-1.65×, cost 2.79×-3.15×, footprint 1.14×-1.47× each update compared state-of-the-art methods. Our codes are available https://github.com/lzhangbv/kfac\_pytorch .